import pathlib

import matplotlib.pyplot
import torch

from checkpoint import Checkpoint
from checkpoint_manager import CheckpointManager
from scaler import Scaler
from simple_pendulum_dataset import SimplePendulumDataset
from simple_pendulum_problem import SimplePendulumProblem
from spinn import Spinn


def main():
    problems = [
        (1, "not_damped", Spinn()),
        (2, "under_damping", Spinn()),
        (3, "over_damping", Spinn(30, 1))
    ]

    for subplot_index, title, model in problems:
        model_path = pathlib.Path(f"./checkpoints/{title}")
        data_path = pathlib.Path(f"./data/{title}.pt")

        dataset = SimplePendulumDataset.load(data_path)
        t_min = float(dataset.t_min())
        t_max = float(dataset.t_max())

        _, checkpoint = CheckpointManager(model_path).last_existed_loaded(Checkpoint.load)
        assert checkpoint is not None
        model.scalar = Scaler(t_min, t_max)
        model.load_state_dict(checkpoint.model_state())
        model.eval()

        original_problem = dataset.problem()
        new_problem = SimplePendulumProblem(
            float(model.g),
            float(model.gamma),
            original_problem.theta0(),
            original_problem.m(),
            original_problem.l_(),
            original_problem.t_step()
        )

        matplotlib.pyplot.subplot(1, len(problems), subplot_index)
        t = torch.linspace(0, 12, 1000)
        matplotlib.pyplot.plot(t, list(original_problem.theta_batch(t)),
                               c="grey", lw=4, alpha=0.6,
                               label="Exact")
        matplotlib.pyplot.plot(t, list(new_problem.theta_batch(t)),
                               c="orange",
                               label="Estimated")
        matplotlib.pyplot.axvspan(t_min, t_max,
                                  color="tab:green", alpha=0.2,
                                  label="Training Region")
        matplotlib.pyplot.plot(t, torch.detach(model(t)),
                               c="tab:green", ls=":",
                               label="Spinn")
        matplotlib.pyplot.legend(loc="upper right")

        matplotlib.pyplot.title(title)
        matplotlib.pyplot.xlabel(r"$t$")
        matplotlib.pyplot.ylabel(r"$\theta$")

    matplotlib.pyplot.subplots_adjust(
        left=0.064, right=0.962, wspace=0.264)
    matplotlib.pyplot.show()


if __name__ == "__main__":
    main()
